import torch
from torch import nn

x = torch.ones(3)
torch.save(x, 'x.pt')

x2 = torch.load('x.pt')
print(x)

# 存储一个Tensor列表
y = torch.zeros(4)
torch.save([x ,y], 'xy.pt')
xy_list = torch.load('xy.pt')
print(xy_list)

# 读取并映射
torch.save({'x': x, 'y': y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
print(xy)

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.hidden = nn.Linear(3, 2)
        self.act = nn.ReLU()
        self.output = nn.Linear(2, 1)

    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)
# 在PyTorch中，Module的可学习参数(即权重和偏差)，模块模型包含在参数中(通过model.parameters()访问)。
# state_dict是一个从参数名称隐射到参数Tensor的字典对象。
net = MLP()
print(net.state_dict())


torch.save(net.state_dict(), "net.pth") # 推荐的文件后缀名是pt或pth
